package com.example.assistant;

import wiremock.com.fasterxml.jackson.databind.JsonNode;
import wiremock.com.fasterxml.jackson.databind.ObjectMapper;

import java.io.BufferedReader;
import java.io.IOException;
import java.io.InputStreamReader;
import java.io.OutputStream;
import java.net.HttpURLConnection;
import java.net.URL;


public class LlmClient {
    private static final String OLLAMA_API_URL = "http://localhost:11434/api/generate";
    private static final ObjectMapper mapper = new ObjectMapper();
    private static final String MODEL_NAME = "qwen2.5:0.5b";

    public static String chatWithModel(String context, String query) throws IOException {
        String prompt = buildPromptWithFewShot(context, query);

        System.out.println("---------------------------\nprompt: " + prompt);
        String escapedPrompt = escapeJson(prompt); // 使用转义方法
        String requestBody = String.format(
                "{\"model\": \"%s\", \"prompt\": \"%s\", \"stream\": false}",
                MODEL_NAME, escapedPrompt);

        HttpURLConnection connection = (HttpURLConnection) new URL(OLLAMA_API_URL).openConnection();
        connection.setRequestMethod("POST");
        connection.setRequestProperty("Content-Type", "application/json");
        connection.setDoOutput(true);

        try (OutputStream os = connection.getOutputStream()) {
            os.write(requestBody.getBytes());
            os.flush();
        }

        if (connection.getResponseCode() != 200) {
            throw new IOException("请求失败: HTTP " + connection.getResponseCode());
        }

        StringBuilder responseText = new StringBuilder();
        try (BufferedReader reader = new BufferedReader(new InputStreamReader(connection.getInputStream()))) {
            String line;
            while ((line = reader.readLine()) != null) {
                if (!line.isEmpty()) {
                    JsonNode responseJson = mapper.readTree(line);
                    String partialResponse = responseJson.get("response").asText();
                    responseText.append(partialResponse);
                }
            }
        }

        return responseText.toString();
    }

    // JSON 特殊字符转义工具方法
    private static String escapeJson(String input) {
        return input.replace("\\", "\\\\")
                .replace("\"", "\\\"")
                .replace("\n", "\\n")
                .replace("\r", "\\r")
                .replace("\t", "\\t");
    }



    /**
     * 构建带有 Few-Shot 示例和格式限制的 Prompt（Java 兼容写法）
     */
    private static String buildPromptWithFewShot(String context, String query) {
        return "你是一个智能家居控制助手，请根据以下设备的标准物模型定义，生成符合要求的 JSON 控制命令。\n"
                + "\n"
                + "请严格按照以下规则执行：\n"
                + "1. 输出必须是合法的 JSON 格式\n"
                + "2. 字段名必须为英文，值必须严格符合物模型定义\n"
                + "3. 不要添加任何解释、多余字段或中文内容\n"
                + "4. 如果无法判断意图，请返回空对象 {}\n"
                + "\n"
                + "以下是设备的标准物模型定义：\n"
                + "\n"
                + context + "\n"
                + "\n"
                + "以下是几个示例输入和输出供你参考：\n"
                + "\n"
                + "示例 1：\n"
                + "用户指令：“打开客厅的灯”\n"
                + "输出：\n"
                + "{\n"
                + "  \"device_type\": \"light\",\n"
                + "  \"property\": \"power\",\n"
                + "  \"value\": \"on\"\n"
                + "}\n"
                + "\n"
                + "示例 2：\n"
                + "用户指令：“把温度调到26度”\n"
                + "输出：\n"
                + "{\n"
                + "  \"device_type\": \"ac\",\n"
                + "  \"property\": \"temperature\",\n"
                + "  \"value\": 26\n"
                + "}\n"
                + "\n"
                + "示例 3：\n"
                + "用户指令：“风扇调到3档”\n"
                + "输出：\n"
                + "{\n"
                + "  \"device_type\": \"fan\",\n"
                + "  \"property\": \"speed\",\n"
                + "  \"value\": 3\n"
                + "}\n"
                + "\n"
                + "示例 4：\n"
                + "用户指令：“加湿器开启自动模式”\n"
                + "输出：\n"
                + "{\n"
                + "  \"device_type\": \"humidifier\",\n"
                + "  \"property\": \"auto_mode\",\n"
                + "  \"value\": \"on\"\n"
                + "}\n"
                + "\n"
                + "现在，请根据以下用户指令生成对应的 JSON 控制命令：\n"
                + "\n"
                + "用户指令：\"" + query + "\"";
    }

}